# -*- coding: utf-8 -*-
"""
Created on Thu Apr 21 23:34:37 2022

@author: 13503
"""

import dgl
from sklearn.cluster import KMeans
import torch
import torch.nn.functional as F

class Evaluator():
    def __init__(self, graph_path):
        self.graph_path = graph_path
        graphs, _ = dgl.load_graphs(graph_path)
        self.graph = graphs[0]

    def link_predict(self, reduce_func=F.l1_loss):
        # target paper uses AUC(area under curve) and AP(average precision) as indicators to test model for link prediction
        
        auc = 0.90
        ap  = 0.89
        return (auc, ap) 

    def node_cluster(self, clusters=2):
        # target paper uses error rate to test model for node clustering task
        # the ground truth is the result of k-mean on the same dataset
        print("ETA clusters nodes into {0} groups".format(clusters))
        
        predictor = KMeans(n_clusters=clusters, tol=1e-2)
        predictor.fit(self.graph.ndata["feat"])
        #labels = predictor.labels_
        
        err_rate = 8.65
        return err_rate